"""
删除数据库表的脚本
"""

from sqlalchemy import inspect, text
from dao.postgresql_connect import engine


def get_all_tables():
    """获取所有数据库表名"""
    inspector = inspect(engine)
    return inspector.get_table_names()


def drop_tables():
    """删除所有数据库表"""
    inspector = inspect(engine)
    tables = inspector.get_table_names()
    
    with engine.connect() as conn:
        for table in tables:
            conn.execute(text(f'DROP TABLE IF EXISTS "{table}" CASCADE'))
            conn.commit()
        print("所有数据库表删除完成")


def drop_specific_table(table_name):
    """删除指定的数据库表
    
    Args:
        table_name: 要删除的表名
    """
    inspector = inspect(engine)
    tables = inspector.get_table_names()
    
    if table_name not in tables:
        print(f"表 {table_name} 不存在")
        return
    
    with engine.connect() as conn:
        conn.execute(text(f'DROP TABLE IF EXISTS "{table_name}" CASCADE'))
        conn.commit()
        print(f"表 {table_name} 删除完成")


if __name__ == '__main__':
    import sys

    if len(sys.argv) > 1:
        drop_specific_table(sys.argv[1])
    else:
        drop_tables()
